package org.oa.ai.model.impl;

import dev.langchain4j.data.message.*;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import org.oa.ai.model.ModelProcessor;

import java.util.List;

public class Qwen2_5Model implements ModelProcessor {

    // 普通返回
    @Override
    public String chatMessage(String message) {
        ChatLanguageModel model = OllamaChatModel.builder()
                .baseUrl("http://localhost:11434")
                .modelName("qwen2.5:14b")
                .build();

        // 添加提示词
        ChatMessage systemMessage = SystemMessage.from("你是一个公文处理助手，你可以帮助我起草公文，续写公文，智能识别公文。必须将数据使用word格式输出。只需要回答我正式内容。");
        UserMessage userMessage = UserMessage.from(
                TextContent.from(message)
        );
        List<ChatMessage> chatMessages = List.of(systemMessage, userMessage);
        // 发送消息
        ChatResponse chat = model.chat(chatMessages);
        return chat.aiMessage().text();
    }

    // 流式返回
    @Override
    public StreamingChatLanguageModel chatMessageStream() {
        StreamingChatLanguageModel model = OllamaStreamingChatModel.builder()
                .baseUrl("http://localhost:11434")
                .modelName("qwen2.5:14b")
                .temperature(0.8)
                .build();
        return model;
    }
}
